#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
资产与IP综合管理系统 - 数据库初始化脚本
用于初始化SQLite数据库和基础数据
"""

import sqlite3
import os
import uuid
from datetime import datetime
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class DatabaseInitializer:
    def __init__(self, db_path="itts.db"):
        """
        初始化数据库初始化器
        
        Args:
            db_path: 数据库文件路径
        """
        self.db_path = db_path
        self.base_dir = os.path.dirname(os.path.abspath(__file__))
        
    def create_database(self):
        """创建数据库和表结构"""
        try:
            # 如果数据库已存在，先备份
            if os.path.exists(self.db_path):
                backup_path = f"{self.db_path}.backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
                os.rename(self.db_path, backup_path)
                logger.info(f"已备份现有数据库到: {backup_path}")
            
            # 连接数据库
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # 启用外键约束
            cursor.execute("PRAGMA foreign_keys = ON")
            
            # 读取并执行建表脚本
            create_script_path = os.path.join(self.base_dir, '01_create_tables_sqlite.sql')
            with open(create_script_path, 'r', encoding='utf-8') as f:
                create_script = f.read()
            
            # 分割并执行SQL语句
            statements = [stmt.strip() for stmt in create_script.split(';') if stmt.strip()]
            for statement in statements:
                if statement:
                    cursor.execute(statement)
            
            conn.commit()
            logger.info("数据库表结构创建完成")
            
            return conn
            
        except Exception as e:
            logger.error(f"创建数据库失败: {e}")
            raise
            
    def init_data(self, conn):
        """初始化基础数据"""
        try:
            cursor = conn.cursor()
            
            # 读取并执行初始化数据脚本
            init_script_path = os.path.join(self.base_dir, '02_init_data_sqlite.sql')
            with open(init_script_path, 'r', encoding='utf-8') as f:
                init_script = f.read()
            
            # 分割并执行SQL语句
            statements = [stmt.strip() for stmt in init_script.split(';') if stmt.strip()]
            for statement in statements:
                if statement and not statement.startswith('--'):
                    cursor.execute(statement)
            
            conn.commit()
            logger.info("基础数据初始化完成")
            
        except Exception as e:
            logger.error(f"初始化数据失败: {e}")
            conn.rollback()
            raise
            
    def verify_database(self, conn):
        """验证数据库创建结果"""
        try:
            cursor = conn.cursor()
            
            # 检查表是否创建成功
            cursor.execute("""
                SELECT name FROM sqlite_master 
                WHERE type='table' AND name NOT LIKE 'sqlite_%'
                ORDER BY name
            """)
            tables = cursor.fetchall()
            
            expected_tables = [
                'DICTIONARY', 'ASSET', 'OUT_OF_BAND', 'OS', 
                'SYSTEM', 'MIDDLEWARE', 'IP', 'ACCOUNT', 'SYSTEM_MIDDLEWARE'
            ]
            
            created_tables = [table[0] for table in tables]
            
            logger.info(f"已创建的表: {', '.join(created_tables)}")
            
            # 检查字典数据
            cursor.execute("SELECT COUNT(*) FROM DICTIONARY")
            dict_count = cursor.fetchone()[0]
            logger.info(f"字典数据条数: {dict_count}")
            
            # 检查示例数据
            cursor.execute("SELECT COUNT(*) FROM ASSET")
            asset_count = cursor.fetchone()[0]
            logger.info(f"设备资产条数: {asset_count}")
            
            cursor.execute("SELECT COUNT(*) FROM IP")
            ip_count = cursor.fetchone()[0]
            logger.info(f"IP地址条数: {ip_count}")
            
            # 验证是否所有预期的表都已创建
            missing_tables = set(expected_tables) - set(created_tables)
            if missing_tables:
                raise Exception(f"缺少表: {', '.join(missing_tables)}")
                
            logger.info("数据库验证通过")
            
        except Exception as e:
            logger.error(f"数据库验证失败: {e}")
            raise
            
    def get_database_info(self, conn):
        """获取数据库信息"""
        try:
            cursor = conn.cursor()
            
            # 获取数据库文件大小
            db_size = os.path.getsize(self.db_path) / 1024 / 1024  # MB
            logger.info(f"数据库文件大小: {db_size:.2f} MB")
            
            # 获取表统计信息
            tables_info = []
            cursor.execute("""
                SELECT name FROM sqlite_master 
                WHERE type='table' AND name NOT LIKE 'sqlite_%'
                ORDER BY name
            """)
            tables = cursor.fetchall()
            
            for table in tables:
                table_name = table[0]
                cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
                count = cursor.fetchone()[0]
                tables_info.append((table_name, count))
            
            logger.info("表数据统计:")
            for table_name, count in tables_info:
                logger.info(f"  {table_name}: {count} 条记录")
                
        except Exception as e:
            logger.error(f"获取数据库信息失败: {e}")
            
    def optimize_database(self, conn):
        """优化数据库"""
        try:
            cursor = conn.cursor()
            
            # 分析表统计信息
            cursor.execute("ANALYZE")
            logger.info("数据库分析完成")
            
            # 清理数据库
            cursor.execute("VACUUM")
            logger.info("数据库清理完成")
            
            conn.commit()
            
        except Exception as e:
            logger.error(f"数据库优化失败: {e}")
            
    def run(self):
        """执行完整的数据库初始化流程"""
        logger.info("开始初始化数据库...")
        
        try:
            # 创建数据库和表结构
            conn = self.create_database()
            
            # 初始化基础数据
            self.init_data(conn)
            
            # 验证数据库
            self.verify_database(conn)
            
            # 获取数据库信息
            self.get_database_info(conn)
            
            # 优化数据库
            self.optimize_database(conn)
            
            logger.info(f"数据库初始化完成: {self.db_path}")
            
        except Exception as e:
            logger.error(f"数据库初始化失败: {e}")
            # 删除可能创建的不完整数据库文件
            if os.path.exists(self.db_path):
                os.remove(self.db_path)
                logger.info("已删除不完整的数据库文件")
            raise
            
        finally:
            if 'conn' in locals():
                conn.close()

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='初始化ITTS数据库')
    parser.add_argument('--db-path', default='itts.db', help='数据库文件路径')
    parser.add_argument('--force', action='store_true', help='强制重新创建数据库')
    
    args = parser.parse_args()
    
    # 检查数据库是否已存在
    if os.path.exists(args.db_path) and not args.force:
        logger.info(f"数据库已存在: {args.db_path}")
        logger.info("如需重新创建，请使用 --force 参数")
        return
    
    # 创建初始化器并运行
    initializer = DatabaseInitializer(args.db_path)
    initializer.run()
    
    logger.info("数据库初始化成功完成！")

if __name__ == "__main__":
    main()